#include "CNN/ckks_mnist_helper.hpp"
#include "intel_itt_wrapper.hpp"

INTEL_ITT_DOMAIN_CREATE(mnist_forward_test, "test.cnn.mnist.forward");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_conv, "HConv");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_sq1, "HSquare1");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_fc1, "HFC1");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_sq2_fc2, "HSquare2-HFC2");

static inline double log2_noise_bound(const helib::Ctxt &ctxt) {
  const double ln2 = NTL::log(NTL::xdouble(2l));
  return NTL::log(ctxt.totalNoiseBound()) / ln2;
}

static inline void Tensor2RawPtxt(std::vector<double> &raw, int row, int col,
                                  const torch::Tensor &a) {
  raw.resize(row * col);
  for (int i = 0; i < row; i++)
    for (int j = 0; j < col; j++) {
      const auto idx = i * col + j;
      raw[idx] = a[i][j].item<float>();
    }
}

static inline void Tensor2RawPtxt_conv(std::vector<double> &raw, int row,
                                       int col, const torch::Tensor &a) {
  raw.resize(row * col);
  for (int i = 0; i < row; i++)
    for (int j = 0; j < col; j++) {
      const auto idx = i * col + j;
      raw[idx] = a[0][0].item<float>();
    }
  // std::cout << mat[0][0] << " " << a[0][0].item() << std::endl;
}

helib::Ctxt ckks_hmnist_client::init_hx(const torch::Tensor &a,
                                        int asz0) const {
  // input data
  std::vector<double> ptxt;
  Tensor2RawPtxt(ptxt, a.size(2), a.size(3), a[asz0][0]);

  // Encryption
  auto ctxt = hmmcc_->encrypt(ptxt, key);
  return std::move(ctxt);
}

helib::Ctxt ckks_hmnist_client::init_hconv(const torch::Tensor &a, int asz0,
                                           int asz1) const {
  // input data
  std::vector<double> ptxt;
  Tensor2RawPtxt_conv(ptxt, 64, 64, a[asz0][asz1]);
  // Encryption
  auto ctxt = hmmcc_->encrypt(ptxt, key);
  return std::move(ctxt);
}

helib::Ctxt ckks_hmnist_client::init_hfc1(const torch::Tensor &a,
                                          int asz0) const {
  // input data
  std::vector<double> ptxt;
  // std::cout << a.size(1) << " " << a.size(2) << std::endl;
  Tensor2RawPtxt(ptxt, a.size(1), a.size(2), a[asz0]);

  // Encryption
  auto ctxt = hmmcc_->encrypt(ptxt, key);
  return std::move(ctxt);
}

helib::Ctxt ckks_hmnist_client::init_hfc2(const torch::Tensor &a) const {
  // input data
  std::vector<double> ptxt;
  Tensor2RawPtxt(ptxt, a.size(0), a.size(1), a);

  // Encryption
  auto ctxt = hmmcc_->encrypt(ptxt, key);
  return std::move(ctxt);
}

ckks_hmnist_client::ckks_hmnist_client(const params<helib::CKKS> &params) {
  auto start_t = omp_get_wtime();
  hmmcc_ = new linearp2_hmmcc(params, PLAINTEXT_NSLOT);
  hmmcc_->report_context(std::cout);
  std::flush(std::cout);
  auto end_t = omp_get_wtime();
  std::cout << "* Setup Helib Context           : " << (end_t - start_t) << " s"
            << std::endl;
}

void ckks_hmnist_client::init_hcnn(encrypted_model<helib::CKKS> &ctxt_model,
                                   const torch::nn::Conv2d &conv1,
                                   const torch::nn::Linear &fc1,
                                   const torch::nn::Linear &fc2) const {
  auto init_hcnn_start_t = omp_get_wtime();
  auto b = init_conv_kernel(conv1->weight.data()); // [49,4,1,1]
  auto c = init_fc1_kernel(fc1->weight.data(), 4); // [4,64,64]
  auto d = init_fc2_kernel(fc2->weight.data());    // [10,64]

// init convolution layer
#pragma omp parallel
#pragma omp single
  {
    auto start_t = omp_get_wtime();
    ctxt_model.conv.resize(b.size(0));
    for (int i = 0; i < b.size(0); i++) {
      for (int j = 0; j < b.size(1); j++) {
        ctxt_model.conv[i].emplace_back(hmmcc_->public_key());
      }
    }
#pragma omp taskloop collapse(2) nogroup
    for (int i = 0; i < b.size(0); i++) {
      for (int j = 0; j < b.size(1); j++) {
        ctxt_model.conv[i][j] = std::move(init_hconv(b, i, j));
      }
    }
    // for (int i = 0; i < b.size(0); i++) {
    //   for (int j = 0; j < b.size(1); j++) {
    //     std::vector<double> res;
    //     hmmcc_->decrypt(res, ctxt_model.conv[i][j]);
    //     auto t = hmmcc_->context()->getP();
    //     std::cout << b[i][j].item<float>() << ", " << res[0] << std::endl;
    //   }
    // }
    auto end_t = omp_get_wtime();
    std::cout << "  * Init convolution layer      : " << (end_t - start_t)
              << " s\n";

    // init fully connected layer1
    start_t = omp_get_wtime();
    for (int i = 0; i < c.size(0); i++) {
      ctxt_model.fc1.emplace_back(hmmcc_->public_key());
    }
#pragma omp taskloop nogroup
    for (int i = 0; i < c.size(0); i++) {
      ctxt_model.fc1[i] = std::move(init_hfc1(c, i));
    }
    end_t = omp_get_wtime();
    std::cout << "  * Init fully connected layer1 : " << (end_t - start_t)
              << " s\n";

    // init fully connected layer2
    start_t = omp_get_wtime();
#pragma omp task
    ctxt_model.fc2 = std::make_shared<helib::Ctxt>(std::move(init_hfc2(d)));
    end_t = omp_get_wtime();
    std::cout << "  * Init fully connected layer2 : " << (end_t - start_t)
              << " s\n";
    auto init_hcnn_end_t = omp_get_wtime();
    std::cout << "* Init HCNN                     : "
              << (init_hcnn_end_t - init_hcnn_start_t) << " s\n";
  }
}

void ckks_hmnist_client::init_input(encrypted_input<helib::CKKS> &ctxt_input,
                                    const torch::Tensor &filter,
                                    const torch::Tensor &x) const {
  auto a = im2matrix(x, filter, 3);
  std::vector<helib::Ctxt> tmp;
  std::vector<std::vector<helib::Ctxt>> tmp2;

#pragma omp parallel
#pragma omp single
  {
    // init input
    auto start_t = omp_get_wtime();
    ctxt_input.data.resize(a.size(0));
    for (int i = 0; i < a.size(0); i++) {
      for (int j = 0; j < a.size(1); j++) {
        ctxt_input.data[i].emplace_back(hmmcc_->public_key());
      }
    }
#pragma omp taskloop collapse(2)
    for (int i = 0; i < a.size(0); i++) {
      for (int j = 0; j < a.size(1); j++) {
        ctxt_input.data[i][j] = std::move(init_hx(a, i));
      }
    }
    auto end_t = omp_get_wtime();
    std::cout << "* Init input                    : " << (end_t - start_t)
              << " s\n";
  }
}

void ckks_hmnist_client::recover_result(
    at::Tensor &out, const encrypted_output<helib::CKKS> &in) const {
  std::vector<double> res;
  hmmcc_->decrypt(res, *in.data);

  for (int i = 0; i < out.size(0); i++) {
    for (int j = 0; j < out.size(1); j++) {
      out.index({i, j}) = res[i * out.size(1) + j];
    }
  }
  out = out.transpose(0, 1);
}

ckks_hmnist_server::ckks_hmnist_server(
    const std::shared_ptr<helib::Context> context,
    const hmm_status<shmm_engine> &status) {
  metric_.hmm_encoding_time -= omp_get_wtime();
  hmme_ = new linear_hmme(status);
  const auto eh_mnk =
      linear_hmme::get_expected_mnk(MATRIX_DIM_M, MATRIX_DIM_N, MATRIX_DIM_K);
  hmme_->register_engine(context->getEA(), eh_mnk);
  metric_.hmm_encoding_time += omp_get_wtime();
}

void ckks_hmnist_server::forward_test(
    encrypted_output<helib::CKKS> &out, const encrypted_input<helib::CKKS> &in,
    const encrypted_model<helib::CKKS> &ctxt_model) const {
  auto &metric = const_cast<struct ckks_metric &>(metric_);

  // Reset metrics
  metric.reset();

  auto ctxt = in.data;
  const auto dim0 = in.data.size();
  const auto dim1 = in.data[0].size();
  metric.forword_test_time -= omp_get_wtime();
  omp_set_nested(1);
#pragma omp parallel
#pragma omp single
  {
    metric.init_err = ctxt[0][0].errorBound();
    metric.init_noise = log2_noise_bound(in.data[0][0]);
    // conv_hnmm
    metric.conv_noise = metric.init_noise;
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_conv); //! ITT - Begin
    metric.hconv_time -= omp_get_wtime();
// #pragma omp parallel for collapse(2)
#pragma omp taskloop collapse(2)
    for (int i = 0; i < dim0; i++)
      for (int j = 0; j < dim1; j++) {
        ctxt[i][j] *= ctxt_model.conv[i][j];
      } //! omp taskloop: Implicit synchronization
    metric.conv_hmul_err = ctxt[0][0].errorBound();

    // reduction
// #pragma omp parallel for
#pragma omp taskloop
    for (int j = 0; j < dim1; j++) {
      for (int i = 1; i < dim0; i++) {
        ctxt[0][j] += ctxt[i][j];
      }
    } //! omp taskloop: Implicit synchronization

    metric.hconv_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metric.conv_noise = log2_noise_bound(ctxt[0][0]) - metric.conv_noise;
    metric.conv_hadd_err = ctxt[0][0].errorBound();

    // sqr1
    metric.sq1_noise = log2_noise_bound(ctxt[0][0]);
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_sq1); //! ITT - Begin
    metric.hsq1_time -= omp_get_wtime();
// #pragma omp parallel for
#pragma omp taskloop
    for (int i = 0; i < dim1; i++) {
      ctxt[0][i].square();
    } //! omp taskloop: Implicit synchronization
    metric.hsq1_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metric.sq1_noise = log2_noise_bound(ctxt[0][0]) - metric.sq1_noise;
    metric.sq1_err = ctxt[0][0].errorBound();

    // hfc1
    metric.fc1_noise = log2_noise_bound(ctxt[0][0]);
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_fc1); //! ITT - Begin
    metric.hfc1_time -= omp_get_wtime();
    // #pragma omp parallel for
    // #pragma omp taskloop
    for (int i = 0; i < dim1; i++) {
      hmme_->cgemm(MATRIX_TYPE::SQUARE, ctxt[0][i], ctxt_model.fc1[i],
                   ctxt[0][i]);
      if (i) {
        ctxt[0][0] += ctxt[0][i];
      }
    } //! omp taskloop: Implicit synchronization
    // for (int i = 1; i < dim1; i++) {
    //   ctxt[0][0] += ctxt[0][i];
    // }
    metric.hfc1_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metric.fc1_noise = log2_noise_bound(ctxt[0][0]) - metric.fc1_noise;
    metric.fc1_err = ctxt[0][0].errorBound();

    // sqr2
    metric.sq2_noise = log2_noise_bound(ctxt[0][0]);
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_sq2_fc2); //! ITT - Begin
    metric.hsq2_time -= omp_get_wtime();
    ctxt[0][0].square();
    metric.hsq2_time += omp_get_wtime();
    metric.sq2_noise = log2_noise_bound(ctxt[0][0]) - metric.sq2_noise;
    metric.sq2_err = ctxt[0][0].errorBound();

    // hfc2
    metric.fc2_noise = log2_noise_bound(ctxt[0][0]);
    metric.hfc2_time -= omp_get_wtime();
    hmme_->cgemm(MATRIX_TYPE::SQUARE, ctxt[0][0], *ctxt_model.fc2, ctxt[0][0]);
    metric.hfc2_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metric.fc2_noise = log2_noise_bound(ctxt[0][0]) - metric.fc2_noise;
    metric.final_noise = log2_noise_bound(ctxt[0][0]);
    metric.left_cap = ctxt[0][0].capacity();
  }
  omp_set_nested(0);

  out.data = std::make_shared<helib::Ctxt>(std::move(ctxt[0][0]));

  metric.err = out.data->errorBound();
  metric.forword_test_time += omp_get_wtime();
}

void ckks_hmnist_server::report_metrics(std::ostream &ss) const {
  // clang-format off
  ss << "================= Cipher CNN =================\n";
  ss << "* HMM Encoding       : " << metric_.hmm_encoding_time << " s\n";
  ss << "* HCNN Time          : " << metric_.forword_test_time << " s\n";
  ss << "  * HConv time       : " << metric_.hconv_time << " s\n";
  ss << "  * HSquare1 time    : " << metric_.hsq1_time << " s\n";
  ss << "  * HFC1 time        : " << metric_.hfc1_time << " s\n";
  ss << "  * HSquare2 time    : " << metric_.hsq2_time << " s\n";
  ss << "  * HFC2 time        : " << metric_.hfc2_time << " s\n";
  ss << "* Final error bound  : " << metric_.err << std::endl;
  ss << "  * Initial error    : " << metric_.init_err << std::endl;
  ss << "  * HConv-HMul error : " << metric_.conv_hmul_err << "(x" << metric_.conv_hmul_err / metric_.init_err  << ")" << std::endl;
  ss << "  * HConv-HAdd error : " << metric_.conv_hadd_err << "(x" << metric_.conv_hadd_err / metric_.conv_hmul_err  << ")" << std::endl;
  ss << "  * HSquare1 error   : " << metric_.sq1_err << "(x" << metric_.sq1_err / metric_.conv_hadd_err  << ")" << std::endl;
  ss << "  * HFC1 error       : " << metric_.fc1_err << "(x" << metric_.fc1_err / metric_.sq1_err  << ")" << std::endl;
  ss << "  * HSquare2 error   : " << metric_.sq2_err << "(x" << metric_.sq2_err / metric_.fc1_err  << ")" << std::endl;
  ss << "  * HFC2 error       : " << metric_.err << "(x" << metric_.err / metric_.sq2_err  << ")" << std::endl;
  ss << "* Final nosie bound  : " << metric_.final_noise  << std::endl;
  ss << "  * Init nosie bound : " << metric_.init_noise << std::endl;
  ss << "  + HConv noise      : " << metric_.conv_noise << std::endl;
  ss << "  + HSquare1 noise   : " << metric_.sq1_noise << std::endl;
  ss << "  + HFC1 noise       : " << metric_.fc1_noise << std::endl;
  ss << "  + HSquare2 noise   : " << metric_.sq2_noise << std::endl;
  ss << "  + HFC2 noise       : " << metric_.fc2_noise << std::endl;
  ss << "* Left capacity      : " << metric_.left_cap << std::endl;
  ss << "==============================================\n";
  // clang-format on
}
